package main

import (
	"bytes"
	"connectrpc.com/connect"
	"fmt"
	"google.golang.org/protobuf/compiler/protogen"
	"google.golang.org/protobuf/types/descriptorpb"
	"google.golang.org/protobuf/types/pluginpb"
	"os"
	"path"
	"path/filepath"
	"strings"
	"unicode/utf8"
)

const (
	contextPackage      = protogen.GoImportPath("context")
	errorsPackage       = protogen.GoImportPath("errors")
	httpPackage         = protogen.GoImportPath("net/http")
	stringsPackage      = protogen.GoImportPath("strings")
	reflectPackage      = protogen.GoImportPath("reflect")
	pathPackage         = protogen.GoImportPath("path")
	connectPackage      = protogen.GoImportPath("connectrpc.com/connect")
	protoreflectPackage = protogen.GoImportPath("google.golang.org/protobuf/reflect/protoreflect")

	jsonPackage                 = protogen.GoImportPath("encoding/json")
	generatedFilenameExtension  = ".connect.pb.go"
	supplementFilenameExtension = ".connect.supplement.pb.go"

	usage                = "See https://connectrpc.com/docs/go/getting-started to learn how to use this plugin.\n\nFlags:\n  -h, --help\tPrint this help and exit.\n      --version\tPrint the version and exit."
	commentWidth         = 97
	protoSyntaxFieldNum  = 12
	protoPackageFieldNum = 2
)

type KeyValue struct {
	Key   string
	Value string
}

func main() {
	if len(os.Args) == 2 && os.Args[1] == "--version" {
		fmt.Fprintln(os.Stdout, connect.Version)
		os.Exit(0)
	}
	if len(os.Args) == 2 && (os.Args[1] == "-h" || os.Args[1] == "--help") {
		fmt.Fprintln(os.Stdout, usage)
		os.Exit(0)
	}
	if len(os.Args) != 1 {
		fmt.Fprintln(os.Stderr, usage)
		os.Exit(1)
	}
	protogen.Options{}.Run(
		func(plugin *protogen.Plugin) error {
			plugin.SupportedFeatures = uint64(pluginpb.CodeGeneratorResponse_FEATURE_PROTO3_OPTIONAL)
			for _, file := range plugin.Files {
				if !file.Generate {
					continue
				}
				if len(file.Services) == 0 {
					return nil
				}
				generatedFilenamePrefixToSlash := filepath.ToSlash(file.GeneratedFilenamePrefix)
				file.GeneratedFilenamePrefix = path.Join(
					path.Dir(generatedFilenamePrefixToSlash),
					string(file.GoPackageName),
					path.Base(generatedFilenamePrefixToSlash),
				)
				generate(plugin, file)
				//补充文件添加
				supplement(plugin, file)
			}
			return nil
		},
	)
}

func generate(plugin *protogen.Plugin, file *protogen.File) {
	generatedFile := plugin.NewGeneratedFile(
		file.GeneratedFilenamePrefix+generatedFilenameExtension,
		file.GoImportPath,
	)
	// 确保目标目录存在
	if err := os.MkdirAll(path.Dir(file.GeneratedFilenamePrefix), 0755); err != nil {
		plugin.Error(err)
		return
	}
	//常量、接口、类型定义
	generatedConstants(generatedFile, file)

	//客户端常量
	generatedClientConstants(generatedFile, file)
	//客户端连接
	generatedNewClient(generatedFile, file)
	for _, service := range file.Services {
		//connect转http的路由handler
		generateHttpServerHandler(generatedFile, service)
	}
}

func generatedNewClient(g *protogen.GeneratedFile, file *protogen.File) {
	for _, service := range file.Services {
		name := changeServiceName(service)
		g.P("func ", name.NewClientFuncName, "(httpClient ", connectPackage.Ident("HTTPClient"), ", baseURL string, opts ...", connectPackage.Ident("ClientOption"), ") ", name.ClientName, " {")
		if len(service.Methods) > 0 {
			g.P("baseURL = ", stringsPackage.Ident("TrimRight"), `(baseURL, "/")`)
		}
		g.P(name.ServiceDescriptor, ` = `,
			g.QualifiedGoIdent(file.GoDescriptorIdent),
			`.Services().ByName("`, service.Desc.Name(), `")`)
		g.P("return &", unexport(name.ClientName), "{")
		for _, method := range service.Methods {
			g.P(unexport(method.GoName), ":", connectPackage.Ident("NewClient"), "[", method.Input.GoIdent, ", ", method.Output.GoIdent, "](")
			g.P("httpClient,")
			g.P("baseURL+" + procedureHttpServiceKeyPath(method) + ",")
			g.P(connectPackage.Ident("WithSchema"), "(", name.ServiceDescriptor, `.Methods().ByName("`+method.GoName+`")),`)
			g.P(connectPackage.Ident("WithClientOptions"), "(opts...),")
			g.P("),")
		}
		g.P("}")
		g.P("}")
		g.P()
	}

}

func generatedClientConstants(g *protogen.GeneratedFile, file *protogen.File) {
	//1. 结构体
	for _, service := range file.Services {
		name := changeServiceName(service)
		g.P("type ", unexport(name.ClientName), " struct {")
		for _, method := range service.Methods {
			g.P(unexport(method.GoName), " *", connectPackage.Ident("Client"), "[", method.Input.GoIdent, ", ", method.Output.GoIdent, "]")
		}
		g.P("}")
	}
	g.P()
	//2. 接口
	for _, service := range file.Services {
		name := changeServiceName(service)
		g.P("type ", name.ClientName, " interface {")
		for _, method := range service.Methods {
			g.P(method.GoName, "(", contextPackage.Ident("Context"), ", *", connectPackage.Ident("Request"), "[", method.Input.GoIdent, "]) (*",
				connectPackage.Ident("Response"), "[", method.Output.GoIdent, "], error)")
		}
		g.P("}")
	}
	g.P()
	//3. 结构体的方法
	for _, service := range file.Services {
		name := changeServiceName(service)
		for _, method := range service.Methods {
			g.P("func (c *", unexport(name.ClientName), ") ", method.GoName,
				"(ctx ", contextPackage.Ident("Context"), ", req *",
				connectPackage.Ident("Request"), "[", method.Input.GoIdent,
				"]) (*", connectPackage.Ident("Response"), "[", method.Output.GoIdent, "], error) {")
			g.P("return c.", unexport(method.GoName), ".CallUnary(ctx, req)")
			g.P("}")
		}
	}
}

var generatedMethods = make(map[string]bool)

func markMethodGenerated(methodName string) {
	generatedMethods[methodName] = true
}
func isMethodGenerated(methodName string) bool {
	_, ok := generatedMethods[methodName]
	return ok
}

// 主函数注册的定义
func supplementRegister(g *protogen.GeneratedFile, file *protogen.File) {
	if len(file.Services) == 0 {
		return
	}
	for _, ser := range file.Services {
		name := changeServiceName(ser)
		g.P("func (s *", Global.HandlerNameType, ") "+name.RegisterName+"(common any) *", Global.HandlerNameType, " {")
		g.P("v := ", reflectPackage.Ident("ValueOf"), "(common).Elem()")
		g.P(Global.FillRouterName, "(v)")
		g.P("return &", Global.AllRouterName)
		g.P("}")
		g.P()
	}
}

// 主函数注册的定义
func generateRegisterOld(g *protogen.GeneratedFile, file *protogen.File) {
	if len(file.Services) == 0 {
		return
	}
	for _, ser := range file.Services {
		name := changeServiceName(ser)
		handlerOption := connectPackage.Ident("HandlerOption")
		g.P("func "+name.RegisterName+"(common any, opts ...", handlerOption, ") ", Global.HandlerNameType, " {")
		for _, service := range file.Services {
			name := changeServiceName(service)
			g.P(name.ServiceDescriptor, ` = `,
				g.QualifiedGoIdent(file.GoDescriptorIdent),
				`.Services().ByName("`, service.Desc.Name(), `")`)
		}
		g.P(Global.AllRouterName, " = make(", Global.HandlerNameType, ",1)")
		g.P("v := ", reflectPackage.Ident("ValueOf"), "(common).Elem()")
		g.P("if v.NumField() == 0 {")
		g.P(Global.FillRouterName, "(v,opts...)")
		g.P("return ", Global.AllRouterName)
		g.P("}")
		g.P("for i := 0; i < v.NumField(); i++ {")
		g.P(Global.FillRouterName, "(v.Field(i),opts...)")
		g.P("}")
		g.P("return ", Global.AllRouterName)
		g.P("}")
		g.P()
	}
}

// 定义每一个函数的方法
func generateHttpServerHandler(g *protogen.GeneratedFile, service *protogen.Service) {
	name := changeServiceName(service)
	for _, method := range service.Methods {
		g.P("func ", procedureProxyMethodName(method), "(srv *", name.UnimplementedServiceName, ", opts ...", connectPackage.Ident("HandlerOption"), ") ", httpPackage.Ident("Handler"), " {")
		isStreamingServer := method.Desc.IsStreamingServer()
		isStreamingClient := method.Desc.IsStreamingClient()
		idempotency := methodIdempotency(method)
		switch {
		case isStreamingClient && !isStreamingServer:
			g.P(procedureHandlerName(method), ` := `, connectPackage.Ident("NewClientStreamHandler"), "(")
		case !isStreamingClient && isStreamingServer:
			g.P(procedureHandlerName(method), ` := `, connectPackage.Ident("NewServerStreamHandler"), "(")
		case isStreamingClient && isStreamingServer:
			g.P(procedureHandlerName(method), ` := `, connectPackage.Ident("NewBidiStreamHandler"), "(")
		default:
			g.P(procedureHandlerName(method), ` := `, connectPackage.Ident("NewUnaryHandler"), "(")
		}
		g.P(procedureHttpServiceKeyPath(method), `,`)
		g.P("srv.", method.GoName, ",")
		g.P(connectPackage.Ident("WithSchema"), "(", name.ServiceDescriptor, `.Methods().ByName("`, method.Desc.Name(), `")`, "),")
		switch idempotency {
		case connect.IdempotencyNoSideEffects:
			g.P(connectPackage.Ident("WithIdempotency"), "(", connectPackage.Ident("IdempotencyNoSideEffects"), "),")
		case connect.IdempotencyIdempotent:
			g.P(connectPackage.Ident("WithIdempotency"), "(", connectPackage.Ident("IdempotencyIdempotent"), "),")
		case connect.IdempotencyUnknown:
		}
		g.P(connectPackage.Ident("WithHandlerOptions"), "(opts...),")
		g.P(")")
		g.P("return ", httpPackage.Ident("HandlerFunc"), `(func(w `, httpPackage.Ident("ResponseWriter"), `, r *`, httpPackage.Ident("Request"), ") {")
		g.P(procedureHandlerName(method), ".ServeHTTP(w, r)")
		g.P("})")
		g.P("}")
		g.P()
	}
	g.P()

}

// 全局变量定义
func generatedConstants(g *protogen.GeneratedFile, file *protogen.File) {
	g.P("package ", file.GoPackageName)
	g.P()
	//1. 常量定义
	g.P("const _ = ", connectPackage.Ident("IsAtLeastVersion1_13_0"))
	g.P()
	if len(file.Services) == 0 {
		return
	}
	g.P("const (")
	for _, service := range file.Services {
		g.P(procedureServiceKeyPath(service), ` = "`, procedureServiceValuePath(service), `"`)
	}
	g.P(")")
	g.P()

	g.P("const (")
	for _, service := range file.Services {
		if len(service.Methods) == 0 {
			return
		}
		for _, method := range service.Methods {
			g.P(procedureHttpServiceKeyPath(method), ` = "`, procedureHttpServiceValuePath(service, method), `"`)
		}
	}
	g.P(")")
	g.P()

	//2. 全局变量定义
	g.P("var (")
	if !isMethodGenerated(Global.AllRouterName) {
		g.P(Global.AllRouterName, " ", Global.HandlerNameType)
		markMethodGenerated(Global.AllRouterName)
	}

	if len(file.Services) == 0 {
		return
	}
	g.P()
	for _, service := range file.Services {
		name := changeServiceName(service)
		g.P(name.ServiceDescriptor, " ", protoreflectPackage.Ident("ServiceDescriptor"))
	}
	g.P()
	for _, service := range file.Services {
		name := changeServiceName(service)
		g.P(service.Comments.Leading.String())
		g.P(name.ServerProxyName, " = map[string]", name.ServerProxyTypeName, "{")
		for _, method := range service.Methods {
			g.P(procedureHttpServiceKeyPath(method), `:`, procedureProxyMethodName(method), `,`)
		}
		g.P("}")
		g.P()
	}
	g.P(")")
	g.P()

	//3. 全局类型定义
	//函数定义
	g.P("type (")
	for _, service := range file.Services {
		name := changeServiceName(service)
		g.P(name.ServerProxyTypeName, " func(srv *", name.UnimplementedServiceName, ", opts ...", connectPackage.Ident("HandlerOption"), ") ", httpPackage.Ident("Handler"))
	}
	g.P(")")
	g.P()

	//结构体定义
	g.P("type (")
	g.P()
	for _, service := range file.Services {
		name := changeServiceName(service)
		g.P(name.UnimplementedServiceName, " struct {")
		g.P("srv ", name.GrpcHttpHandler)
		g.P("}")
		g.P()
	}
	g.P()
	g.P(")")

	g.P()
	if !isMethodGenerated(Global.HttpFilter) {
		g.P("type ", Global.HttpFilter, " struct {")
		g.P("WhiteList []string")
		g.P("ServerName string")
		g.P("Func  ", Global.MiddleFunc)
		g.P("}")
		markMethodGenerated(Global.HttpFilter)
	}
	g.P()

	g.P()
	if !isMethodGenerated(Global.MiddleFunc) {
		g.P("type ", Global.MiddleFunc, " func(w ", httpPackage.Ident("ResponseWriter"), ", r *", httpPackage.Ident("Request"), ") error")
		markMethodGenerated(Global.MiddleFunc)
	}
	g.P()

	g.P()
	if !isMethodGenerated(Global.HandlerConf) {
		g.P("type ", Global.HandlerConf, " struct {")
		g.P("HttpServer ", httpPackage.Ident("Handler"))
		g.P("Filter ", Global.HttpFilter)
		g.P("}")
		markMethodGenerated(Global.HandlerConf)
	}
	g.P()

	g.P()
	if !isMethodGenerated(Global.HandlerNameType) {
		g.P("type ", Global.HandlerNameType, " map[string]", Global.HandlerConf)
		markMethodGenerated(Global.HandlerNameType)
	}
	g.P()

	//定义接口类型
	for _, service := range file.Services {
		name := changeServiceName(service)
		g.P("type ", name.GrpcHttpHandler, " interface {")
		for _, method := range service.Methods {
			g.AnnotateSymbol(name.ServerProxyName+"."+method.GoName, protogen.Annotation{Location: method.Location})
			signatureHanler(g, method, false)
		}
		g.P("}")
		g.P()
	}
}

func signatureHanler(g *protogen.GeneratedFile, method *protogen.Method, needConnect bool) {
	ctxName := "ctx "
	reqName := "req "
	streamName := "stream "
	if method.Desc.IsStreamingClient() && method.Desc.IsStreamingServer() {
		// bidi streaming
		g.P(method.GoName, "("+ctxName, " ", g.QualifiedGoIdent(contextPackage.Ident("Context")), ", "+
			streamName+reqName+" *BidiStream"+
			g.QualifiedGoIdent(method.Input.GoIdent)+", "+g.QualifiedGoIdent(method.Output.GoIdent)+
			") error")
	}
	if method.Desc.IsStreamingClient() {
		// client streaming
		g.P(method.GoName, "("+ctxName, " ", g.QualifiedGoIdent(contextPackage.Ident("Context"))+", "+
			streamName, reqName, " *"+g.QualifiedGoIdent(method.Input.GoIdent)+
			") (*"+g.QualifiedGoIdent(method.Output.GoIdent)+",error)")
	}
	if method.Desc.IsStreamingServer() {
		// server streaming
		g.P(method.GoName, "("+ctxName, " ", g.QualifiedGoIdent(contextPackage.Ident("Context"))+
			", "+reqName, " *",
			g.QualifiedGoIdent(method.Input.GoIdent)+", "+
				streamName+"*"+
				g.QualifiedGoIdent(method.Output.GoIdent)+""+
				") error")
	}
	g.P(method.GoName, "("+ctxName, " ", g.QualifiedGoIdent(contextPackage.Ident("Context")),
		", "+reqName, " *", g.QualifiedGoIdent(method.Input.GoIdent), ") (*", g.QualifiedGoIdent(method.Output.GoIdent), ", error)")
}

// 方法代理函数名称定义key
func procedureProxyMethodName(method *protogen.Method) string {
	return fmt.Sprintf("_%sHandler", method.Desc.Name())
}

// http路由路径定义key
func procedureHttpServiceKeyPath(method *protogen.Method) string {
	return fmt.Sprintf("%s%sProcedure", method.Parent.GoName, method.GoName)
}

// http路由路径定义value
func procedureHttpServiceValuePath(service *protogen.Service, method *protogen.Method) string {
	return fmt.Sprintf("/%s/%s", service.Desc.FullName(), method.Desc.Name())
}

// 服务路径定义key
func procedureServiceKeyPath(service *protogen.Service) string {
	return fmt.Sprintf("%sName", service.Desc.Name())
}

// 服务路径定义value
func procedureServiceValuePath(service *protogen.Service) string {
	return fmt.Sprintf("/%s/", service.Desc.FullName())
}

type GlobalCollection struct {
	HandlerNameType string
	AllRouterName   string
	FillRouterName  string
	StartName       string
	HandlerConf     string
	MiddleFunc      string
	HttpFilter      string
	HandlerFilter   string
}

var Global = GlobalCollection{
	HandlerNameType: "HandlerMap",
	AllRouterName:   "allRouter",
	FillRouterName:  "fillRouter",
	StartName:       "StartServer",
	HandlerConf:     "HandlerConf",
	MiddleFunc:      "MiddleFunc",
	HttpFilter:      "HttpFilter",
	HandlerFilter:   "HandlerFilter",
}

type GlobalServiceCollection struct {
	Base                     string
	UnimplementedServiceName string
	ServerProxyName          string
	ServerProxyTypeName      string
	GrpcHttpHandler          string
	ServiceDescriptor        string
	ClientName               string
	NewClientFuncName        string
	RegisterName             string
}

func changeServiceName(service *protogen.Service) GlobalServiceCollection {
	base := service.GoName
	return GlobalServiceCollection{
		Base:                     base,
		UnimplementedServiceName: fmt.Sprintf("Unimplemented%sServiceHandler", base),
		ServerProxyName:          fmt.Sprintf("%sServiceProxy", base),
		ServerProxyTypeName:      fmt.Sprintf("%sServerHandler", base),
		GrpcHttpHandler:          fmt.Sprintf("%sServiceGrpcHttpHandler", base),
		ServiceDescriptor:        fmt.Sprintf("%sServiceDescriptor", base),
		ClientName:               fmt.Sprintf("%sServiceClient", base),
		NewClientFuncName:        fmt.Sprintf("New%sClient", base),
		RegisterName:             fmt.Sprintf("New%sRegister", base),
	}
}
func methodIdempotency(method *protogen.Method) connect.IdempotencyLevel {
	methodOptions, ok := method.Desc.Options().(*descriptorpb.MethodOptions)
	if !ok {
		return connect.IdempotencyUnknown
	}
	switch methodOptions.GetIdempotencyLevel() {
	case descriptorpb.MethodOptions_NO_SIDE_EFFECTS:
		return connect.IdempotencyNoSideEffects
	case descriptorpb.MethodOptions_IDEMPOTENT:
		return connect.IdempotencyIdempotent
	case descriptorpb.MethodOptions_IDEMPOTENCY_UNKNOWN:
		return connect.IdempotencyUnknown
	}
	return connect.IdempotencyUnknown
}
func procedureHandlerName(m *protogen.Method) string {
	return fmt.Sprintf("%s%sHandler", unexport(m.Parent.GoName), m.GoName)
}
func unexport(s string) string {
	lowercased := strings.ToLower(s[:1]) + s[1:]
	switch lowercased {
	// https://go.dev/ref/spec#Keywords
	case "break", "default", "func", "interface", "select",
		"case", "defer", "go", "map", "struct",
		"chan", "else", "goto", "package", "switch",
		"const", "fallthrough", "if", "range", "type",
		"continue", "for", "import", "return", "var":
		return "_" + lowercased
	default:
		return lowercased
	}
}

func wrapComments(g *protogen.GeneratedFile, elems ...any) {
	text := &bytes.Buffer{}
	for _, el := range elems {
		switch el := el.(type) {
		case protogen.GoIdent:
			fmt.Fprint(text, g.QualifiedGoIdent(el))
		default:
			fmt.Fprint(text, el)
		}
	}
	words := strings.Fields(text.String())
	text.Reset()
	var pos int
	for _, word := range words {
		numRunes := utf8.RuneCountInString(word)
		if pos > 0 && pos+numRunes+1 > commentWidth {
			g.P("// ", text.String())
			text.Reset()
			pos = 0
		}
		if pos > 0 {
			text.WriteRune(' ')
			pos++
		}
		text.WriteString(word)
		pos += numRunes
	}
	if text.Len() > 0 {
		g.P("// ", text.String())
	}
}
